using System;
using System.Collections.Generic;
using System.Linq;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics;
using Atomic.Structures;

namespace Atomic.Thermodynamics.PartitionFunctions
{
	public partial class IsobaricPartitionFunction
	{
		[Serializable]
		public partial class UnivariateIsobaricPartitionFunction : IGibbsGrandPotential
		{
			private Atom atom;
			private double[] counts;
			private double count_min, count_max;
			private RollingCache<Tuple<double, double>, InnerSolver> solvers;

			public UnivariateIsobaricPartitionFunction(IEnumerable<IIsobaricPartitionFunctionEntry> entries)
			{
				if (entries == null)
				{
					throw new ArgumentNullException();
				}

				Entries = new ReadOnlyList<IIsobaricPartitionFunctionEntry>(entries);

				if (Entries.Count == 0)
				{
					throw new ArgumentException();
				}

				Atom[] atoms = new ReadOnlyList<Atom>(entries.SelectMany(e => e.Composition.Atoms).Distinct().Where(a => a != Atom.Vacancy)).ToArray();

				if (atoms.Length != 1)
				{
					throw new ArgumentException("Only 1-atom atom systems are supported.");
				}

				atom = atoms[0];
				counts = entries.Select(p => p.Composition.Count(atom)).ToArray();
				count_min = counts.Min();
				count_max = counts.Max();

				// Rolling cache of solvers for different temperatures and pressures tuples.
				solvers = new RollingCache<Tuple<double, double>, InnerSolver>(1000);
			}

			public double FreeEnergy(double temperature, double pressure, Composition composition)
			{
				double x;
				if (!Count(composition, out x))
				{
					return double.NaN;
				}

				return CreateSolver(temperature, pressure).GibbsPotential(x);
			}

			public double Volume(double temperature, double pressure, Composition composition)
			{
				double x;
				Count(composition, out x);

				return CreateSolver(temperature, pressure).Volume(x);
			}

			public ChemicalPotential ChemicalPotential(double temperature, double pressure, Composition composition)
			{
				double x;
				if (!Count(composition, out x))
				{
					return new ChemicalPotential(atom, double.NaN);
				}

				double mu = CreateSolver(temperature, pressure).ChemicalPotential(x);
				return new ChemicalPotential(atom, mu);
			}

			public double VolumetricThermalExpansionCoefficient(double temperature, double pressure, Composition composition)
			{
				double h = 0.0001;

				//double r = EquilibriumWignerSeitzRadius(temperature);
				//double rh = EquilibriumWignerSeitzRadius(temperature + h);
				//double alpha = 1.0 / r * (rh - r) / h;
				double v = Volume(temperature, pressure, composition);
				double vh = Volume(temperature + h, pressure, composition);
				double alpha = 1.0 / v * (vh - v) / h;

				// Linear thermal expansion.
				//alpha /= 3.0;

				return alpha;

				// Account for change in configurational probabilities as well as explicit expansion for each configuration.
				//throw new System.NotImplementedException();
			}
			
			private bool Count(Composition composition, out double count)
			{
				if (composition == null)
				{
					throw new ArgumentNullException();
				}

				foreach (Atom a in composition.Atoms)
				{
					if (a == Atom.Vacancy || a == Atom.Undefined)
					{
						// Ignore vacancies.
						continue;
					}

					if (a != atom)
					{
						throw new ArgumentException();
					}
				}

				count = composition.Count(atom);

				// Must be inside the convex hull spanned by the composition of the entries.
				return count >= count_min && count <= count_max;
			}

			private InnerSolver CreateSolver(double temperature, double pressure)
			{
				if (temperature < 0.0)
				{
					throw new ArgumentOutOfRangeException();
				}

				InnerSolver solver;
				if (!solvers.TryGetValue(Tuple.Create<double, double>(temperature, pressure), out solver))
				{
					solver = new InnerSolver(this, temperature, pressure);
					solvers.Add(Tuple.Create<double, double>(temperature, pressure), solver);
				}

				return solver;
			}

			public ReadOnlyList<IIsobaricPartitionFunctionEntry> Entries
			{
				get;
				private set;
			}

			[Serializable]
			private class InnerSolver
			{
				private UnivariateIsobaricPartitionFunction parent;
				private int n;
				private double[] energies, volumes;
				private double beta, tol, mu_min, mu_max;
				private RollingCache<double, double> mus;

				public InnerSolver(UnivariateIsobaricPartitionFunction parent, double t, double p)
				{
					this.parent = parent;

					n = parent.Entries.Count;

					// Precompute free energies for the given fixed temperature and pressure.
					energies = new double[n];
					volumes = new double[n];
					for (int i = 0; i < n; i++)
					{
						energies[i] = parent.Entries[i].Potential.FreeEnergy(t, p);
						volumes[i] = parent.Entries[i].Potential.Volume(t, p);
					}

					double kb = 8.617332478e-5; // eV/K
					beta = 1.0 / (kb * t);

					// Rolling cache of chemical potentials.
					mus = new RollingCache<double, double>(10000);

					// Determine range of chemical potentials.
					tol = 1.0e-8;
					mu_min = SolveChemicalPotential(parent.count_min + tol, -1.0e8, 1.0e8);
					mu_max = SolveChemicalPotential(parent.count_max - tol, -1.0e8, 1.0e8);
				}

				private double LogPartitionFunction(double mu, out double[] f)
				{
					double[] e = new double[n];
					double e0 = double.NegativeInfinity;
					for (int i = 0; i < n; i++)
					{
						e[i] = beta * mu * parent.counts[i] - beta * energies[i];
						e0 = Math.Max(e0, e[i]);
					}

					// Compute partition function with prefactor introduced for numerical stability.
					double z = 0.0;
					f = new double[n];
					for (int i = 0; i < n; i++)
					{
						f[i] = Math.Exp(e[i] - e0) * parent.Entries[i].Degeneracy;
						z += f[i];
					}

					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						f[i] /= z;
					}

					// Divide out prefactor again after taking the logarithm.
					double lz = Math.Log(z) + e0;

					return lz;
				}

				public double GrandPotential(double mu)
				{
					double[] f;
					double lz = LogPartitionFunction(mu, out f);

					double phi = -lz / beta;

					return phi;
				}

				public double Count(double mu)
				{
					double[] f;
					LogPartitionFunction(mu, out f);

					double x = 0.0;
					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						x += f[i] * parent.counts[i];
					}

					return x;
				}

				private double SolveChemicalPotential(double x, double mu_min, double mu_max)
				{
					double mu;
					if (!mus.TryGetValue(x, out mu))
					{
						mu = RootFinder.Bisection(y => Count(y) - x, mu_min, mu_max);
						mus.Add(x, mu);
					}

					return mu;
				}

				public double ChemicalPotential(double x)
				{
					return SolveChemicalPotential(x, mu_min, mu_max);
				}
				
				public double GibbsPotential(double x)
				{
					// Legendre transform of the grand-potential.
					double mu = ChemicalPotential(x);
					double phi = GrandPotential(mu);
					double g = phi + mu * x;

					return g;
				}

				public double Volume(double x)
				{
					double mu = ChemicalPotential(x);

					double[] f;
					LogPartitionFunction(mu, out f);

					double v = 0.0;
					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						v += f[i] * volumes[i];
					}

					return v;
				}
			}

			IHelmholtzGrandPotential IGibbsGrandPotential.HelmholtzPotential
			{
				get
				{
					throw new NotImplementedException();
				}
			}
		}
	}
}
